"""
Image Retriever Module
图像检索模块
"""

import os
import pickle
import logging
from typing import List, Tuple, Dict
import torch
import faiss
import numpy as np
from PIL import Image
from transformers import CLIPProcessor, CLIPModel

logger = logging.getLogger(__name__)


class ImageRetriever:
    """图像检索器"""
    
    def __init__(self, config: dict):
        """
        初始化图像检索器
        
        Args:
            config: 配置字典
        """
        self.config = config
        self.clip_model = None
        self.processor = None
        self.faiss_index = None
        self.image_id_to_captions = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        self._load_models()
        self._load_knowledge_base()
    
    def _load_models(self):
        """加载CLIP模型和处理器"""
        try:
            model_path = self.config['models']['clip_model_path']
            self.clip_model = CLIPModel.from_pretrained(model_path)
            self.processor = CLIPProcessor.from_pretrained(model_path)
            self.clip_model.to(self.device)
            self.clip_model.eval()
            logger.info(f"CLIP模型加载成功，使用设备: {self.device}")
        except Exception as e:
            logger.error(f"CLIP模型加载失败: {e}")
            raise
    
    def _load_knowledge_base(self):
        """加载FAISS索引和映射关系"""
        try:
            # 加载FAISS索引
            knowledge_base_path = self.config['knowledge_base']['knowledge_base_path']
            if not os.path.exists(knowledge_base_path):
                raise FileNotFoundError(f"知识库文件不存在: {knowledge_base_path}")
            
            self.faiss_index = faiss.read_index(knowledge_base_path)
            
            # 加载映射关系
            captions_path = self.config['knowledge_base']['image_id_to_captions_path']
            if not os.path.exists(captions_path):
                raise FileNotFoundError(f"映射文件不存在: {captions_path}")
            
            with open(captions_path, 'rb') as f:
                self.image_id_to_captions = pickle.load(f)
            
            logger.info(f"知识库加载成功，包含 {self.faiss_index.ntotal} 个向量")
            
        except Exception as e:
            logger.error(f"加载知识库失败: {e}")
            raise
    
    def extract_features(self, image: Image.Image) -> np.ndarray:
        """
        提取单张图像的特征向量
        
        Args:
            image: PIL图像对象
            
        Returns:
            np.ndarray: 特征向量
        """
        try:
            inputs = self.processor(images=image, return_tensors="pt")
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            with torch.no_grad():
                image_features = self.clip_model.get_image_features(**inputs)
                image_features = image_features / image_features.norm(dim=-1, keepdim=True)
                return image_features.cpu().numpy().astype('float32')
                
        except Exception as e:
            logger.error(f"提取图像特征失败: {e}")
            raise
    
    def retrieve_similar_images(self, query_image: Image.Image, top_k: int = None) -> List[Tuple[int, float]]:
        """
        检索最相似的top_k张图像
        
        Args:
            query_image: 查询图像
            top_k: 返回的相似图像数量
            
        Returns:
            List[Tuple[int, float]]: 图像ID和相似度分数的列表
        """
        try:
            if top_k is None:
                top_k = self.config['retrieval']['top_k']
            
            # 提取查询图像特征
            query_features = self.extract_features(query_image)
            
            # 执行检索
            scores, indices = self.faiss_index.search(query_features, top_k)
            
            # 获取图像ID列表
            image_ids = list(self.image_id_to_captions.keys())
            results = []
            
            for score, idx in zip(scores[0], indices[0]):
                if idx < len(image_ids):
                    image_id = image_ids[idx]
                    results.append((image_id, float(score)))
            
            logger.info(f"检索到 {len(results)} 个相似图像")
            return results
            
        except Exception as e:
            logger.error(f"图像检索失败: {e}")
            raise
    
    def get_retrieved_captions(self, query_image: Image.Image) -> List[List[str]]:
        """
        获取检索到的图像的描述列表
        
        Args:
            query_image: 查询图像
            
        Returns:
            List[List[str]]: 每个相似图像对应的描述列表
        """
        try:
            # 检索相似图像
            similar_images = self.retrieve_similar_images(query_image)
            
            # 获取对应的描述
            captions_list = []
            for image_id, score in similar_images:
                if image_id in self.image_id_to_captions:
                    captions = self.image_id_to_captions[image_id]
                    captions_list.append(captions)
                else:
                    logger.warning(f"图像ID {image_id} 没有找到对应的描述")
                    captions_list.append([])
            
            return captions_list
            
        except Exception as e:
            logger.error(f"获取检索描述失败: {e}")
            raise
